#include <init.h>
#include "LogFactorialTable.h"

using namespace std;
using namespace clUtil;

unsigned int gNumClasses = 3;
unsigned int gPrecursorTolerance = 1500;
unsigned int gPeakTolerance = 500;
float gTICCutoff = 0.98f;
bool gPPMTolerance = false;
bool gPrecursorPPMTolerance = false;
unsigned int gBestScoreListSize = 5;
unsigned int gMaxPrecursorCharge = 3;
unsigned int gMaxMissedCleavages = UINT_MAX;
unsigned int gMaxPeakCount = 301;
unsigned int gMinTerminiCleavages = 2;
char* gConfigFile = NULL;
char* gWorkDir = NULL;
int gCPUS = -1;
char* gProteinDatabase = NULL;
vector<char*> gMSMSDataFiles;
unsigned int gBufferSize = 1024 * 1024 * 1024; //1GB

const char* gVersion = "0.9";

//Performance auto-tuning section
vector<AlgorithmConfig> gDeviceAlgorithmConfig;

void usage()
{
  printf("Usage: specmaster [options] --ProteinDatabase <FASTA database>"
         " <MSMS files>\n");
  exit(1);
}

void parseArgs( int argc, char** argv )
{
  int c;
  int option_index = 0;

  while(1)
  {
    static struct option long_options[]=
    {
      {"cfg",                       required_argument, 0, 'c'},
      {"workdir",                   required_argument, 0, 'w'},
      {"cpus",                      required_argument, 0, 'p'},
      {"ProteinDatabase",           required_argument, 0, 'P'},
      {"help",                      no_argument      , 0, 'h'},
      {"NumIntensityClasses",       required_argument, 0, 'i'},
      {"AvePrecursorMzTolerance",   required_argument, 0, 'm'},
      {"FragmentMzTolerance",       required_argument, 0, 'f'},
      {"TicCutoffPercentage",       required_argument, 0, 't'},
      {"MaxResultRank",             required_argument, 0, 'r'},
      {"NumChargeStates",           required_argument, 0, 'C'},
      {"MaxMissedCleavages",        required_argument, 0, 'M'},
      {"MaxPeakCount",              required_argument, 0, 'a'},
      {"MinTerminiCleavages",       required_argument, 0, 'T'},
      {"BufferSize",                required_argument, 0, 'b'},
      {0,0,0,0}
    };

    c = getopt_long(argc, 
                    argv, 
                    "c:w:p:P:hi:m:f:t:r:C:M:a:T:b:", 
                    long_options, 
                    &option_index);

    if(c == -1)
    {
      break;
    }

    switch(c)
    {
      case 'c':
        gConfigFile = optarg;
        break;
      case 'w':
        gWorkDir = optarg;
        break;
      case 'p':
        if(sscanf(optarg, "%d", &gCPUS) != 1)
        {
          fprintf(stderr, 
                  "Warning: ignoring value for command line option "
                  "cpus.\n" );
        }
        break;
      case 'P':
        gProteinDatabase = optarg;
        break;
      case 'h':
        usage();
        break;
      case 'i':
        if(sscanf(optarg, "%d", &gNumClasses) != 1)
        {
          fprintf(stderr, 
                  "Warning: ingnoring value for command line option "
                  "NumIntensityClasses.\n");
        }
        break;
      case 'm':
        if(sscanf(optarg, "%d", &gPrecursorTolerance) != 1)
        {
          fprintf(stderr, 
                  "Warning: ignoring value for command line option "
                  "AvePrecursorMzTolerance.\n");
        }
        break;
      case 'f':
        if(sscanf(optarg, "%d", &gPeakTolerance ) != 1)
        {
          fprintf(stderr, 
                  "Warning: ignoring value for command line option "
                  "FragmentMzTolerance.\n");
        }
        break;
      case 't':
        if(sscanf(optarg, "%f", &gTICCutoff) != 1)
        {
          fprintf(stderr, 
                  "Warning: ignoring value for command line option "
                  "TicCutoffPercentage.\n");
        }
        break;
      case 'r':
        if(sscanf(optarg, "%d", &gBestScoreListSize) != 1)
        {
          fprintf(stderr, 
                  "Warning: ignoring value for command line option "
                  "MaxResultRank.\n");
        }
        break;
      case 'C':
        if(sscanf(optarg, "%d", &gMaxPrecursorCharge) != 1)
        {
          fprintf(stderr, 
                  "Warning: ignoring value for command line option "
                  "NumChargeStates.\n");
        }
        break;
      case 'M':
        if(sscanf(optarg, "%d", &gMaxMissedCleavages) != 1)
        {
          fprintf(stderr, 
                  "Warning: ignoring value for command line option "
                  "MaxMissesCleavages.\n");
        }
        break;
      case 'a':
        if(sscanf(optarg, "%d", &gMaxPeakCount) != 1)
        {
          fprintf(stderr, 
                  "Warning: ingnoring value for command line option "
                  "MaxPeakCount.\n");
        }
        break;
      case 'T':
        if(sscanf(optarg, "%d", &gMinTerminiCleavages) != 1)
        {
          fprintf(stderr, 
                  "Warning: ignoring value for command line option "
                  "MinTerminiCleavages.\n");
        }
        break;
      case 'b':
        if(sscanf(optarg, "%d", &gBufferSize) != 1)
        {
          fprintf(stderr, 
                  "Warning: ignoring value for command line option "
                  "BufferSize.\n");
        }
        else
        {
          gBufferSize *= 1024 * 1024;
        }
        break;
      case '?':
        usage();
        break;
      default:
        abort();
    }
  }

  gMSMSDataFiles.clear();

  for(unsigned int curMSMS = (unsigned int)optind; 
      curMSMS < (unsigned int)argc; 
      curMSMS++)
  {
    gMSMSDataFiles.push_back(argv[curMSMS]);
  }
}

//Return Monoisotopic mass in milliDaltons
static unsigned int getAcidMonoMass(unsigned char acid)
{
  switch(acid)
  {
    case 'A':
      return 71037;
    case 'C':
      return 103009;
    case 'D':
      return 115027;
    case 'E':
      return 129043;
    case 'F':
      return 147068;
    case 'G':
      return 57021;
    case 'H':
      return 137059;
    case 'I':
      return 113084;
    case 'K':
      return 128095;
    case 'L':
      return 113084;
    case 'M':
      return 131040;
    case 'N':
      return 114043;
    case 'P':
      return 97053;
    case 'Q':
      return 128059;
    case 'R':
      return 156101;
    case 'S':
      return 87032;
    case 'T':
      return 101048;
    case 'V':
      return 99068;
    case 'W':
      return 186079;
    case 'Y':
      return 163063;
    default:
      return 0;
  }
}

void tuneDevices()
{
  size_t numDevices = Device::GetDevices().size();

  gDeviceAlgorithmConfig.resize(numDevices);

  //For each device, figure out if we should use global, local, or image memory
  for(size_t curDeviceID = 0; curDeviceID < numDevices; curDeviceID++)
  {
    const Device& curDevice = Device::GetDevices()[curDeviceID];
    const DeviceInfo& curDeviceInfo = curDevice.getDeviceInfo();
    
    if(curDeviceInfo.Type == CL_DEVICE_TYPE_CPU)
    {
      gDeviceAlgorithmConfig[curDeviceID].LocalWorkSize = 1;
      gDeviceAlgorithmConfig[curDeviceID].SpectrumStorage = 
        kStorageMechanismGlobalMemory;
    }
    else
    {
      if(curDeviceInfo.Name == "Tahiti")
      {
        gDeviceAlgorithmConfig[curDeviceID].LocalWorkSize = 64;
        gDeviceAlgorithmConfig[curDeviceID].SpectrumStorage = 
          kStorageMechanismGlobalMemory;
      }
      else
      {
        gDeviceAlgorithmConfig[curDeviceID].LocalWorkSize = 64;
        gDeviceAlgorithmConfig[curDeviceID].SpectrumStorage = 
          kStorageMechanismLocalMemory;
      }
    }
    
  }
}

void createHeader()
{
  FILE* header;

  header = fopen("kernels/tables.cl", "w");

  if(header == NULL)
  {
    printf("Could not create configuration header file!\n");
    exit(1);
  }

  //Write AA table
  fprintf(header, "#pragma once\n");
  fprintf(header, "__constant unsigned int gMonoMass[256] = \n  {\n");

  for(unsigned char curAA = 0; curAA < 255; curAA++)
  {
    fprintf(header, "    %u,\n", getAcidMonoMass(curAA));
  }

  fprintf(header, "    %u\n  };\n\n", getAcidMonoMass(255));

  //Write log factorial table
  fprintf(header, 
          "__constant float gLogFact[%u] = \n  {\n", 
          (unsigned int)gNumFactorials);

  for(size_t curFact = 0; curFact < gNumFactorials - 1; curFact++)
  {
    fprintf(header, "    %Ef,\n", gLogFact[curFact]);
  }

  fprintf(header, "    %Ef\n  };\n\n", gLogFact[gNumFactorials - 1]);

  fclose(header);

  //Write constants.cl
  header = fopen("kernels/constants.cl", "w");

  fprintf(header, "#ifndef CONSTANTS_CL\n");
  fprintf(header, "#define CONSTANTS_CL\n");
  fprintf(header, "#define kMassH 1008u\n");
  fprintf(header, "#define kMassOH 17007u\n");
  fprintf(header, "#define kMassH2O 18015u\n");
  fprintf(header, "#define kNumClasses %uu\n", gNumClasses);
  fprintf(header, "#define kPeakTolerance %uu\n", gPeakTolerance);
  fprintf(header, "#define kBestScoreListSize %uu\n", gBestScoreListSize);
  fprintf(header, "#define kMaxPrecursorCharge %uu\n", gMaxPrecursorCharge);
  fprintf(header, "#define kTICCutoffFactor %f\n", gTICCutoff);
  fprintf(header, "#define kMaxPeakCount %uu\n", gMaxPeakCount);

  if(gPrecursorPPMTolerance == true)
  {
    fprintf(header, "#define RELATIVE_PRECURSOR_TOLERANCE\n");
    fprintf(header, 
            "#define kPrecursorTolerance %uu\n", 
            gPrecursorTolerance);
  }
  else
  {
    fprintf(header, 
            "#define kPrecursorTolerance %uu\n", 
            gPrecursorTolerance);
  }
  if(gPPMTolerance == true)
  {
    fprintf(header, "#define RELATIVE_TOLERANCE\n");
  }
  else
  {
    fprintf(header, 
            "#define kFragmentMZTolerance %ff\n", 
            (float)gPeakTolerance / 1e3);
  }

  //Generate autotuning variables
  size_t numDevices = Device::GetDevices().size();

  for(size_t curDeviceID = 0; curDeviceID < numDevices; curDeviceID++)
  {
    fprintf(header, "#if CLUTIL_DEVICE_ID == %ld\n", curDeviceID);

    switch(gDeviceAlgorithmConfig[curDeviceID].SpectrumStorage)
    {
      case kStorageMechanismImage:
        fprintf(header, "#define IMAGE_SPECTRUM\n");
        break;
      case kStorageMechanismGlobalMemory:
        fprintf(header, "#define GLOBAL_MEMORY_SPECTRUM\n");
        break;
      case kStorageMechanismLocalMemory:
        fprintf(header, "#define LOCAL_MEMORY_SPECTRUM\n");
        break;
      default:
        break;
    }

    fprintf(header, 
            "#define kNumThreadsPerBlock %uu\n", 
            (unsigned int)gDeviceAlgorithmConfig[curDeviceID].LocalWorkSize);

    fprintf(header, "#endif\n");
  }
  
  fprintf(header, "#endif\n");

  fclose(header);
}

//TODO: Clean this function up
void parseConfig()
{
  if(gConfigFile == NULL)
  {
    return;
  }

  string line;
  ifstream configFile(gConfigFile);

  while(getline(configFile, line))
  {
    string option;
    istringstream sstream(line);

    sstream.exceptions(ios::failbit);

    sstream >> option;

    if(option == "NumIntensityClasses")
    {
      try 
      {
        sstream >> gNumClasses;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option " 
             << option << endl;
      }
    }
    else if(option == "AvePrecursorMzTolerance")
    {
      try
      {
        sstream >> gPrecursorTolerance;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option "
             << option << endl;
      }
    }
    else if(option == "FragmentMzTolerance")
    {
      try
      {
        sstream >> gPeakTolerance;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option "
             << option << endl;
      }
    }
    else if(option == "TicCutoffPercentage")
    {
      try
      {
        sstream >> gTICCutoff;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option "
             << option << endl;
      }
    }
    else if(option == "MaxResultRank")
    {
      try
      {
        sstream >> gBestScoreListSize;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option "
             << option << endl;
      }
    }
    else if(option == "NumChargeStates")
    {
      try
      {
        sstream >> gMaxPrecursorCharge;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option "
             << option << endl;
      }
    }
    else if(option == "MaxMissedCleavages")
    {
      try
      {
        sstream >> gMaxMissedCleavages;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option "
             << option << endl;
      }
    }
    else if(option == "MaxPeakCount")
    {
      try
      {
        sstream >> gMaxPeakCount;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option "
             << option << endl;
      }
    }
    else if(option == "MinTerminiCleavages")
    {
      try
      {
        sstream >> gMinTerminiCleavages;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option "
             << option << endl;
      }
    }
    else if(option == "BufferSize")
    {
      try
      {
        sstream >> gBufferSize;
      }
      catch(ios_base::failure e)
      {
        cout << "Warning: ignoring value for config file option "
             << option << endl;
      }
    }
    else 
    {
      cout << "Warning: ignoring unrecognized config option "
           << option << endl;
    }
  }
}

void init(int argc, char** argv)
{
  parseArgs(argc, argv);
  parseConfig();
  parseArgs(argc, argv);

  if(gProteinDatabase == NULL || gMSMSDataFiles.size() == 0)
  {
    usage();
  }

  Device::FetchDevices();

  tuneDevices();

  createHeader();
}
